package cn.jly.flink.source2sink.datastream.source;

import cn.jly.flink.entity.SensorReading;
import cn.jly.flink.utils.FlinkUtils;
import org.apache.commons.lang3.RandomUtils;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @PackageName cn.jly.flink.source2sink.datastream.source
 * @ClassName CustomSource
 * @Description 自定义数据源
 * @Author 姬岚洋
 * @Date 2021/1/26 下午 8:17
 */
public class CustomSource {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = FlinkUtils.getStreamExecutionEnv();

        env.addSource(new CustomSensorReadingSourceFunction("mySource"))
                .print();

        FlinkUtils.executeStream(env, "CustomSource");
    }

    public static class CustomSensorReadingSourceFunction implements SourceFunction<SensorReading> {
        // 标识符，控制数据生产
        private boolean running = true;

        private final String sourceName;

        // 定义初始10个传感器
        private final HashMap<String, Double> idAndTempMap = new HashMap<>(10);

        {
            // 设置初始温度值，60度~89度
            for (int i = 0; i < 10; i++) {
                idAndTempMap.put("sensor_" + (i + 1), 60d + RandomUtils.nextDouble(0d, 30d));
            }
        }

        public CustomSensorReadingSourceFunction(String sourceName) {
            this.sourceName = sourceName;
        }


        @Override
        public void run(SourceContext<SensorReading> ctx) throws Exception {
            System.out.println("数据源名称 = " + sourceName);

            while (running) {
                // 模拟传感器温度数据值
                for (Map.Entry<String, Double> entry : idAndTempMap.entrySet()) {
                    final SensorReading sensorReading = new SensorReading();
                    // 温度随机波动
                    sensorReading.setId(entry.getKey());
                    // 模拟温度升高
                    if (RandomUtils.nextBoolean()) {
                        sensorReading.setTemperature(entry.getValue() + RandomUtils.nextDouble(0, 3));
                    } else {
                        // 模拟温度降低
                        sensorReading.setTemperature(entry.getValue() - RandomUtils.nextDouble(0, 3));
                    }
                    sensorReading.setTimestamp(System.currentTimeMillis() + RandomUtils.nextInt(0, 3));

                    ctx.collect(sensorReading);
                }

                // 控制生成频率
                TimeUnit.SECONDS.sleep(1);
            }
        }

        @Override
        public void cancel() {
            this.running = false;
        }

        public String getSourceName() {
            return sourceName;
        }
    }
}
